open Sast
open Ast
open Type
open Operator

module SymTbl = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)

(* Find symbols from scope *)
let rec findfun id = function
    | [] ->  
        if Builtin.exists id 
            then Builtin.get_ret_type id
            else failwith ("Undeclared Function:" ^ id)
    | loc::scp ->
        let (vars,funs) = loc in
        if SymTbl.mem id funs then
        SymTbl.find id funs
        else
        findfun id scp
;;

let rec findvar id = function
    | [] -> failwith ("Undeclared Variable:" ^ id)
    | loc::scp ->
        let (vars,funs) = loc in
        if SymTbl.mem id vars 
            then SymTbl.find id vars
            else if SymTbl.mem id funs
                then SymTbl.find id funs
                else findvar id scp
;;


let chk ast = 
    let rec expr env = function
              Ast.LInt e -> Sast.LInt e, Int 
            | Ast.LFloat e -> Sast.LFloat e, Float
            | Ast.Binop(e1, op, e2) ->
                            let e1 = expr env e1
                            and e2 = expr env e2 in
                            let _, t1 = e1
                            and _, t2 = e2
                            in (match op with
                            |  Add|Sub -> (match (t1,t2) with 
                            | (Int, _) | (RInt, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                                    | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                                    | (_, _)                    -> failwith "Unsupported Add/Sub")
                            | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Float
                                    | (_, _)                    -> failwith "Unsupported Add/Sub")
                            | (Vector a, Vector b) -> (match a with
                                | b -> Sast.Binop(e1, op, e2), Vector a)
                            | (_, _) -> failwith "Unsupported Add/Sub")

                            | Equal|Neq|Greater|Geq|Leq|Less -> 
                                    (* use Int instead of boolean *)
                                (match (t1,t2) with 
                                     | (Int, _) | (RInt, _)| (Float, _) | (RFloat, _) -> (match (t1,t2) with
                                         | (_, Int) | (_, RInt)| (_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Int
                                         | (_, _) -> failwith "Unsupported Compare")
                                     | (_, _) -> failwith "Unsupported Compare")

                            | Mult -> (match (t1, t2) with
                                | (Int, _) | (RInt, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                                    | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                                    | (_, Vector a)               -> Sast.Binop(e1, op, e2), Vector a
                                    | (_, _)                    -> failwith "Unsupported Mult")
                                | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Float
                                    | (_, _)                    -> failwith "Unsupported Mult")
                                    | (_, Vector a)               -> Sast.Binop(e1, op, e2), Vector a
                                | (_, _) -> failwith "Unsupported Mult")

                            | Div -> (match (t1, t2) with
                                | (Int, _) | (RInt, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                                    | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                                    | (_, _)                    -> failwith "Unsupported Div")
                                | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Float
                                    | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                                    | (_, _)                    -> failwith "Unsupported Div")
                                | (_, _) -> failwith "Unsupported Div")

                            | Rmndr -> (match (t1, t2) with
                                | (Int, _) | (RInt, _) -> (match (t1,t2) with
                                    | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                                    | (_, _)                    -> failwith "Unsupported Reminder")
                                | (_, _) -> failwith "Unsupported Reminder")

                            | Innrp -> (match (t1,t2) with 
                                | (Vector a, Vector b)  -> Sast.Binop(e1, op, e2), Float
                                | (_, _) -> failwith "Unsupported Innrp")
                            )

            | Ast.Uop(op, e) -> 
                            let e = expr env e in
                            let _, t = e in
                            (match op with
                            | Minus -> (match t with
                            | Int|Float|RInt|RFloat|(Vector _) -> Sast.Uop(op, e), t
                            | _ -> failwith "Unsupported Minus")
                            
                            (* Right now, Not could be operated on any type
                             * Everything could be regarded as boolean 1 expect 0*) 
                            | Not -> Sast.Uop(op, e), t)
           | Ast.LString str -> Sast.LString str, Str
           | Ast.LList e -> 
                           Sast.LList (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), List
           | Ast.LVctr e -> Sast.LVctr (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), Vector (List.length e)
           | Ast.LTple e -> Sast.LTple (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), Tuple
           | Ast.Elmt (s, e) ->
                           let e = expr env e in Sast.Elmt (s, e), Int
           | Ast.Id e -> Sast.Id e, (findvar e env)
           | Ast.Call (s, e) ->  Sast.Call(s, List.fold_left(fun ls ex -> ls@[(expr env ex)]) [] e), (findfun s env)
    in

    let rec stmt env = function
             Ast.Expr(e) -> Sast.Expr(expr env e), env

           | Ast.If(e, s1, s2) ->
                    let e = expr env e in
                    let _, t = e in
                    (match t with
                    | Int -> Sast.If(e, fst (stmt env s1), fst (stmt env s2)), env
                    | _ -> failwith "Predicate of if must be integer")

           | Ast.VDclr(t, str, s) -> (match env with
                [] -> failwith "empty scope in FDclr"
                | (lvars, lfuns)::globals -> 
                    let lvars = SymTbl.add str t lvars in 
                    let env = (lvars, lfuns)::globals in (
                Sast.VDclr(t, str, fst (stmt env s)), env))
           | Ast.Block b -> Sast.Block (fst (List.fold_left (fun (st, env') b -> let stmt' = (stmt env' b) in ((fst stmt')::st, (snd stmt'))) ([], env) b)), env
           | Ast.While (s, e) -> failwith "While is not supported!"
           | Ast.Scope (s1, s2) -> 
                    let env = (SymTbl.empty, SymTbl.empty)::env 
                    in
                    let (s1' , env)= 
                        List.fold_left 
                            (fun (st, env') s1 -> 
                                let stmt' = (stmt env' s1)
                                in
                                ((fst stmt')::st, (snd stmt'))
                            ) 
                            ([], env) s1
                    in Sast.Scope (s1',  fst (stmt env s2)), env
           | Ast.FDclr (t1, str, args, s) -> (match env with
                [] -> failwith "empty scope in FDclr"
                | (lvars, lfuns)::globals -> 
                    let lfuns = SymTbl.add str t1 lfuns in 
                    let lvars_new = List.fold_left (fun scope' (typ, var) -> (SymTbl.add var typ scope')) SymTbl.empty args in
                    let env = [(lvars_new, SymTbl.empty); (lvars, lfuns)]@globals in (
                Sast.FDclr(t1, str, args, fst(stmt env s)), env))
    in

    let emptyScope = [(SymTbl.empty, SymTbl.empty)] in

    fst (List.fold_left (fun (prog, env) st -> let stmt' = (stmt env st) in prog@[(fst stmt')], snd(stmt')) ([], emptyScope) ast)

